/*	
*	File: ipv4-header.h
*	Date: 2011/11/05
*	User: Ronmaru
*	Desctiption: ip v4 header class
*/

#ifndef IPV4_HEADER_HPP
#define IPV4_HEADER_HPP

#include <algorithm>
#include <boost/asio/ip/address_v4.hpp>

// Packet header for IPv4.
//
// The wire format of an IPv4 header is:
// 
// 0               8               16                             31
// +-------+-------+---------------+------------------------------+      ---
// |       |       |               |                              |       ^
// |version|header |    type of    |    total length in bytes     |       |
// |  (4)  | length|    service    |                              |       |
// +-------+-------+---------------+-+-+-+------------------------+       |
// |                               | | | |                        |       |
// |        identification         |0|D|M|    fragment offset     |       |
// |                               | |F|F|                        |       |
// +---------------+---------------+-+-+-+------------------------+       |
// |               |               |                              |       |
// | time to live  |   protocol    |       header checksum        |   20 bytes
// |               |               |                              |       |
// +---------------+---------------+------------------------------+       |
// |                                                              |       |
// |                      source IPv4 address                     |       |
// |                                                              |       |
// +--------------------------------------------------------------+       |
// |                                                              |       |
// |                   destination IPv4 address                   |       |
// |                                                              |       v
// +--------------------------------------------------------------+      ---
// |                                                              |       ^
// |                                                              |       |
// /                        options (if any)                      /    0 - 40
// /                                                              /     bytes
// |                                                              |       |
// |                                                              |       v
// +--------------------------------------------------------------+      ---

const int IPV4_MAX_HEADER_LENGTH = 60;

class IPV4header
{
public:
  IPV4header() { std::fill(rep_, rep_ + sizeof(rep_), 0); }
//   IPV4header(std::string &init_str) { 
//     for (int i = 0; i < init_str.length(); i += 2) {
//       rep_[i / 2] = Encode(GetDigitValue(init_str[i]), GetDigitValue(init_str[i + 1]));
//     }
//   }

  void SetVersion(int n) { rep_[0] &= 0x0F; rep_[0] |= (n << 4); }
  void SetHeaderLength(int n) { rep_[0] &= 0xF0; rep_[0] |= n; }
  void SetTypeOfService(int n) { rep_[1] = n; }
  void SetTotalLength(int n) { Encode(2, 3, n); }
  void SetIdentification(int n) { Encode(4, 5, n); }
  void SetDontFragment(int n) { rep_[6] &= 0xBF; rep_[6] |= n << 6; }
  void SetMoreFragments(int n) { rep_[6] &= 0xDB; rep_[6] |= n << 5; }
  void SetFragmentOffset(int n) { Encode(6, 7, (n & 0x1FFF) | (Decode(6, 7) & 0xE000)); }
  void SetTimeToLive(int n) { rep_[8] = n; }
  void SetProtocol(int n) { rep_[9] = n; }
  void SetHeaderChecksum(int n) { Encode(10, 11, n); }
  void SetSourceAddress(boost::asio::ip::address_v4 address) 
  {
    boost::asio::ip::address_v4::bytes_type bytes = address.to_bytes();
    rep_[12] = bytes[0];
    rep_[13] = bytes[1];
    rep_[14] = bytes[2];
    rep_[15] = bytes[3];
  }
  void SetDestinationAddress(boost::asio::ip::address_v4 address)
  {
    boost::asio::ip::address_v4::bytes_type bytes = address.to_bytes();
    rep_[16] = bytes[0];
    rep_[17] = bytes[1];
    rep_[18] = bytes[2];
    rep_[19] = bytes[3];
  }

  unsigned char GetVersion() const { return (rep_[0] >> 4) & 0xF; }
  unsigned short GetHeaderLength() const { return (rep_[0] & 0xF) * 4; }
  unsigned char GetTypeOfService() const { return rep_[1]; }
  unsigned short GetTotalLength() const { return Decode(2, 3); }
  unsigned short GetIdentification() const { return Decode(4, 5); }
  bool GetDontFragment() const { return (rep_[6] & 0x40) != 0; }
  bool GetMoreFragments() const { return (rep_[6] & 0x20) != 0; }
  unsigned short GetFragmentOffset() const { return Decode(6, 7) & 0x1FFF; }
  unsigned int GetTimeToLive() const { return rep_[8]; }
  unsigned char GetProtocol() const { return rep_[9]; }
  unsigned short GetHeaderChecksum() const { return Decode(10, 11); }

  boost::asio::ip::address_v4 GetSourceAddress() const
  {
    boost::asio::ip::address_v4::bytes_type bytes
      = { { rep_[12], rep_[13], rep_[14], rep_[15] } };
    return boost::asio::ip::address_v4(bytes);
  }

  boost::asio::ip::address_v4 GetDestinationAddress() const
  {
    boost::asio::ip::address_v4::bytes_type bytes
      = { { rep_[16], rep_[17], rep_[18], rep_[19] } };
    return boost::asio::ip::address_v4(bytes);
  }

  friend std::ostream& operator<<(std::ostream& os, const IPV4header& header)
  { return os.write(reinterpret_cast<const char*>(header.rep_), header.GetHeaderLength()); }

  friend std::istream& operator>>(std::istream& is, IPV4header& header)
  {
    is.read(reinterpret_cast<char*>(header.rep_), 20);
    if (header.GetVersion() != 4)
      is.setstate(std::ios::failbit);
    std::streamsize options_length = header.GetHeaderLength() - 20;
    if (options_length < 0 || options_length > 40)
      is.setstate(std::ios::failbit);
    else
      is.read(reinterpret_cast<char*>(header.rep_) + 20, options_length);
    return is;
  }

  void ComputeChecksum() 
  {
    SetHeaderChecksum(0);
    unsigned int sum = 0;

    for (int i = 0; i < IPV4_MAX_HEADER_LENGTH; i += 2) {
      sum += rep_[i] << 8;
      sum += rep_[i+1];
    }

    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    SetHeaderChecksum(~sum);
  }

private:
  unsigned short Decode(int a, int b) const
  { return (rep_[a] << 8) + rep_[b]; }

  void Encode(int a, int b, unsigned short n)
  {
    rep_[a] = static_cast<unsigned char>(n >> 8);
    rep_[b] = static_cast<unsigned char>(n & 0xFF);
  }

//   unsigned char Encode(int a, int b) const
//   {
//     return static_cast<unsigned char>((a << 4) + b);
//   }
// 
//   unsigned int GetDigitValue(char digit) const 
//   {
//     if (digit >= '0' && digit <= '9')
//     {
//       return digit - '0';
//     }
//     else if (digit >= 'a' && digit <= 'f')
//     {
//       return 10 + digit - 'a';
//     }
//     return -1;
//   }

  unsigned char rep_[IPV4_MAX_HEADER_LENGTH];
};

#endif // IPV4_HEADER_HPP